# -*- coding: utf-8 -*-
# created on 2018/05/15
from mathsolver.functions.shulie.ShuLie_Transform import DengChaBiEqNormalization
from mathsolver.functions.base import *
from mathsolver.functions.base.base import new_latex
import math
from mathsolver.functions.mathematica.mathematicaSolve import MathematicaSolve, MathematicaSimplify
from mathsolver.functions.shulie.Fundamental import *
from sympy.abc import q
from mathsolver.functions.shulie.basic import sl_search_multiple


class ShuLieYiBanGetPoly(BaseFunction):
    def solver(self, *args):
        seq = args[0].sympify()
        new_seqs = args[1].sympify()
        seq_names = sl_shulie_name(seq)
        seq_xiangs = sl_xiangs(seq)
        seq_values = []
        if len(seq_names) == 2:
            sl = self.search(str(seq_names[0]))
            a_name = sl.aName
            n_name = sl.nName
            if not sl.an:
                sl.get_TongXiangGongShi()
                self.steps += sl.steps_TongXiang
                self.label.update(sl.label_TongXiang)
            tong_xiangs = sl.an
            if not sl.sn:
                sl.get_nSums()
                self.steps += sl.steps_nSums
                self.label.update(sl.label_nSums)
            n_sums = sl.sn
            for tong_xiang, nSum in zip(tong_xiangs, n_sums):
                new_seq = seq
                for seq_xiang in seq_xiangs:
                    new_item = sympify(seq_xiang)
                    new_item_name = new_item.func
                    new_item_index = new_item.args[0]
                    if str(a_name) == str(new_item_name):
                        new_item_value = tong_xiang.subs({n_name: new_item_index})
                        new_seq = new_seq.subs({new_item: new_item_value})
                        new_seq = new_seq.expand().simplify()
                    else:
                        new_item_value = n_sums.subs({n_name: new_item_index})
                        new_seq = new_seq.subs({new_item: new_item_value})
                        new_seq = new_seq.expand().simplify()
                seq_values.append(new_seq)
            seq_values = set(seq_values)
            seq_values = list(seq_values)
            self.steps.append(["", "∴ %s" % (",".join(["%s = %s" % (new_latex(seq), new_latex(value)) for value in seq_values]))])
            self.output.append(BaseNumbers(seq_values))
            self.label.add("求数列的通项公式")
            self.label.add("求数列的前n项和")
        else:
            seq_name = seq_names[0]
            sl = self.search(str(seq_name))
            sl.eqs_update(new_seqs)
            sl.get_shouxiang()
            sl.get_nxiangs()
            sl.get_dituigongshi()
            sl.get_xiangs()
            sl.get_nineqs()
            a_name = sl.aName
            a_name = sympify(a_name)
            if str(seq_name) == str(a_name):
                if not sl.an:
                    sl.get_tongxianggongshi()
                    self.steps += sl.steps_TongXiang
                    self.label.update(sl.label_TongXiang)
                tong_xiangs = sl.an
                flag = []
                for seq_xiang in seq_xiangs:
                    new_item = sympify(seq_xiang)
                    new_item_index = new_item.args[0]
                    if sympify('n') in new_item_index.free_symbols:
                        flag.append(1)
                    else:
                        flag.append(0)
                if len(tong_xiangs) == 1 and isinstance(tong_xiangs[0], BasePieceShuLie) and sum(flag) > 0:
                    tong_xiang = tong_xiangs[0].sympify()
                    new_tong_xiang = []
                    for item in tong_xiang:
                        new_seq = seq
                        _, value, _range = item
                        for seq_xiang in seq_xiangs:
                            new_item = sympify(seq_xiang)
                            new_item_index = new_item.args[0]
                            new_item_value = value.subs({sl.nName: new_item_index})
                            new_seq = new_seq.subs({new_item: new_item_value})
                            if len(str(new_seq)) > 15:
                                new_seq = MathematicaSimplify().solver(BasePoly(new_seq)).output[0].sympify()
                        new_tong_xiang.append([seq, new_seq, _range])
                    new_an = BasePieceShuLie({"var": sl.nName, "name": a_name(sl.nName),
                                              "expression": new_tong_xiang})
                    self.output.append(new_an)
                    self.steps.append(["", "∴ %s " % (new_an.printing())])
                else:
                    for tong_xiang in tong_xiangs:
                        new_seq = seq
                        for seq_xiang in seq_xiangs:
                            new_item = sympify(seq_xiang)
                            new_item_index = new_item.args[0]
                            if isinstance(tong_xiang, BasePieceShuLie):
                                new_tong_xiang = tong_xiang.expression
                                for item in new_tong_xiang:
                                    n_range = item[1]
                                    if new_item_index in n_range:
                                        new_item_value = item[0].subs({sl.nName: new_item_index})
                                        break
                            else:
                                new_item_value = tong_xiang.subs({sl.nName: new_item_index})
                            new_seq = new_seq.subs({new_item: new_item_value})
                            new_seq = new_seq.expand().simplify()
                            new_seq_children = get_all_child(new_seq, lambda tt: isinstance(tt, Pow) and tt.args[1] - sympify('1/2') == 0)
                            if new_seq_children:
                                for new_seq_item in new_seq_children:
                                    new_seq_item_base = new_seq_item.args[0]
                                    if isinstance(new_seq_item_base, Pow) and new_seq_item_base.args[1] == 2:
                                        new_seq_item_value = new_seq_item_base.args[0]
                                        new_seq = new_seq.subs({new_seq_item: new_seq_item_value})
                        seq_values.append(new_seq)
                    seq_values = set(seq_values)
                    seq_values = list(seq_values)
                    self.steps.append(["", "∴ %s" % (",".join(["%s = %s" % (new_latex(seq), new_latex(value)) for value in seq_values]))])
                    self.output.append(BaseNumbers(seq_values))
                self.label.add("求数列的通项公式")
            else:
                if not sl.an:
                    sl.get_tongxianggongshi()
                    self.steps += sl.steps_TongXiang
                    self.label.update(sl.label_TongXiang)
                if not sl.sn:
                    sl.get_nsums()
                    self.steps += sl.steps_nSums
                    self.label.update(sl.label_nSums)
                n_sums = sl.sn
                for nSum in n_sums:
                    new_seq = seq
                    sns_expression = []
                    for seq_xiang in seq_xiangs:
                        new_item = sympify(seq_xiang)
                        new_item_index = new_item.args[0]
                        if isinstance(nSum, BasePieceShuLie):
                            sns = nSum.sympify()
                            for sn in sns:
                                sn_name, sn_value, sn_cond = sn
                                if sn_cond == "n为奇数":
                                    if not new_item_index.free_symbols:
                                        if new_item_index % 2 == 1:
                                            new_item_value = sn_value.subs({sl.nName: new_item_index})
                                            new_seq = new_seq.subs({new_item: new_item_value})
                                            new_seq = new_seq.expand().simplify()
                                            break
                                    else:
                                        sns_expression.append([sn_value, "n为奇数"])
                                elif sn_cond == "n为偶数":
                                    if not new_item_index.free_symbols:
                                        if new_item_index % 2 == 0:
                                            new_item_value = sn_value.subs({sl.nName: new_item_index})
                                            new_seq = new_seq.subs({new_item: new_item_value})
                                            new_seq = new_seq.expand().simplify()
                                            break
                                    else:
                                        sns_expression.append([sn_value, "n为偶数"])
                        else:
                            new_item_value = nSum.subs({sl.nName: new_item_index})
                            new_seq = new_seq.subs({new_item: new_item_value})
                            new_seq = new_seq.expand().simplify()
                    if not sns_expression:
                        seq_values.append(new_seq)
                if sns_expression:
                    new_n_sums = BasePieceShuLie({"var": sympify('n'), "name": sympify('S_(n)'),
                                                  "expression": sns_expression})
                    self.output.append(new_n_sums)
                else:
                    seq_values = set(seq_values)
                    seq_values = list(seq_values)
                    self.steps.append(["", "∴ %s" % (",".join(["%s = %s" % (new_latex(seq), new_latex(value)) for value in seq_values]))])
                    self.output.append(BaseNumbers(seq_values))
                self.label.add("求数列的前n项和")
        return self


class ShuLieAnExprNormalization(BaseFunction):
    """
    a_(?.n)化为S_(?.n)
    """
    def solver(self, *args):
        expr = args[0].sympify()
        s_name = args[1]
        p = r'[a-z]_\(.*?\)'
        m = re.findall(p, str(expr))
        m = set(m)
        m = list(m)
        if m:
            for item in m:
                new_item = sympify(item)
                new_item_index = new_item.args[0]
                if new_item_index == 1:
                    new_item_value = s_name(new_item_index)
                else:
                    new_item_value = s_name(new_item_index) - s_name(new_item_index - 1)
                expr = expr.subs({new_item: new_item_value})
        self.output.append(BasePoly(expr))
        return self


class ShuLiePeriodValuesUpdate(BaseFunction):
    """
    求周期数列的元素值
    """
    def solver(self, *args):
        sl = args[0]
        period = args[1]
        target_values = args[2]
        xiangs = []
        for eq in sl.a1:
            xiangs.append(eq)
        for eq in sl.Xiangs_Eqs:
            xiangs.append(eq)
        condition_eqs = []
        for eq in xiangs:
            expr = eq[0] - eq[1]
            expr_xiangs = sl_xiangs(expr)
            if expr_xiangs:
                for expr_xiang in expr_xiangs:
                    new_item = sympify(expr_xiang)
                    new_item_name = new_item.func
                    new_item_index = new_item.args[0]
                    if str(new_item_name) == str(sl.aName):
                        t = sympify('t')
                        ineqs = [[new_item_index + period * t, '>=', 1], [new_item_index + period * t, '<=', period]]
                        solutions = MathematicaSolve().solver(BaseIneqs(ineqs)).output[0].sympify()
                        t_value = solutions[t]
                        tleft, tright, tlopen, tropen = t_value.args
                        tleft = int(math.ceil(tleft))
                        tright = int(math.floor(tright))
                        t_value = range(tleft, tright + 1)
                        assert len(t_value) == 1
                        t_value = t_value[0]
                        try:
                            new_item_value = target_values[new_item_index - t_value * period - 1]
                        except Exception:
                            new_item_value = target_values[new_item_index + t_value * period - 1]
                        expr = expr.subs({new_item: new_item_value})
                        self.label.add("求数列的通项公式")
                    else:
                        pass
            condition_eqs.append([expr, S.Zero])
        if len(condition_eqs) == 1:
            stepsolver = MathematicaSolve().solver(BaseEq(condition_eqs[0]))
        else:
            stepsolver = MathematicaSolve().solver(BaseIneqs(condition_eqs))
        solutions = stepsolver.output[0].sympify()
        new_target_values = []
        for target_value in target_values:
            target_value_answers = sl_search_multiple(solutions, target_value)
            assert len(target_value_answers) == 1
            target_value_answer = target_value_answers[0]
            target_value = target_value.subs(target_value_answer)
            new_target_values.append(target_value)
        self.output.append(BaseNumbers(new_target_values))
        return self


class ShuLieNXiangsEqsTransformation(BaseFunction):
    def solver(self, *args):
        n_xiangs_eqs = args[0]
        condition_eqs = []
        for eq in n_xiangs_eqs:
            if len(eq) == 2:
                expr = eq[0] - eq[1]
            else:
                expr = eq[0] - eq[2]
            from mathsolver.functions.shulie.ShuLie_getSeqPoly import ShuLieGetSeqPoly
            new_known = dict(self.known)
            stepsolver = ShuLieGetSeqPoly(new_known).solver(BaseSeqPoly(expr))
            self.label.update(stepsolver.label)
            self.steps += stepsolver.steps
            new_exprs = stepsolver.output[0].sympify()
            for new_expr in new_exprs:
                if len(eq) == 2:
                    condition_eqs.append([new_expr, S.Zero])
                else:
                    left, op, right = eq
                    condition_eqs.append([new_expr, op, S.Zero])
        self.output.append(BaseIneqs(condition_eqs))
        return self


class ShuLieDengChaBiSolveEqs(BaseFunction):
    def solver(self, *args):
        new_seqs = args[0].sympify()
        condition_eqs = []
        for eq in new_seqs:
            if len(eq) == 2:
                new_eq1 = [eq[0], eq[1]]
                stepsolver = DengChaBiEqNormalization(self.known).solver([new_eq1])
                new_eqs1 = stepsolver.output[0].sympify()
                for new_eq1 in new_eqs1:
                    if new_eq1 not in condition_eqs:
                        condition_eqs.append(new_eq1)
                self.label.update(stepsolver.label)
            elif len(eq) == 3:
                stepsolver = DengChaBiEqNormalization(self.known).solver([eq])
                self.label.update(stepsolver.label)
                new_eqs = stepsolver.output[0].sympify()
                condition_eqs.extend(new_eqs)
        if len(condition_eqs) == 1:
            self.steps.append(["", "%s" % (BaseEq(condition_eqs[0]).printing())])
            stepsolver = MathematicaSolve().solver(BaseEq(condition_eqs[0]))
            self.label.add("解方程")
        elif len(condition_eqs) > 1:
            self.steps.append(["", "%s" % (BaseIneqs(condition_eqs).printing())])
            stepsolver = MathematicaSolve().solver(BaseIneqs(condition_eqs), args[1])
            self.label.add("解方程组")
        else:
            raise Exception("to do")
        solutions = stepsolver.output[0]
        self.steps.append(["", "解得: %s" % (solutions.printing())])
        self.output.append(solutions)
        return self
